import torch

from raw.get_model import get_model

class BaseDefense:
    def __init__(self, name, model_name, config, retrain, ignore_output = False):
        self.name = name
        self.config = config
        self.retrain = retrain
        self.ignore_output = ignore_output
        self.model = get_model(model_name, self.get_model_options())

    # 修改网络结构的防御需要覆写此方法，返回要传给模型的模型选项（dict）
    def get_model_options(self):
        return {}

    # 重新训练模型的防御需要覆写此方法，返回防御好的网络参数（state_dict）
    def defend(self):
        print('For defenses that retrain the model, method `defend` must be implemented')
        raise NotImplementedError

    # 对输入图像进行变形的防御需要覆写此方法（参数和返回值均为CPU上的torch.Tensor）
    def transform_input(self, images):
        return images

    # 修改预测过程的防御需要覆写此方法（参数和返回值与原始网络输入输出相同，且均为模型所在设备上的torch.Tensor）
    def predict(self, images):
        return self.model(images)

    # 测试
    def test(self, loader):
        self.model.eval() # 将网络设置为测试模式

        total = 0
        corr = 0
        with torch.no_grad():
            for images, labels in loader:
                images = self.transform_input(images) # 对输入图像进行变形

                images = images.to(self.model.device)
                labels = labels.to(self.model.device)

                output = self.predict(images) # 将图像传入网络
                pred = torch.argmax(output, dim = 1) #找出概率最大的标签

                total += labels.size(0)
                corr += (pred == labels).sum().item()

        ratio = corr / total  #分类正确比率

        print('Test: Accuracy: {}/{} ({:.0f}%)'.format(
            corr, total, ratio * 100))

        return ratio

    # 对于需要重新训练的防御，加载防御后的网络参数；对于不需要重新训练的防御，加载最初训练的网络参数
    def load_enhanced_state(self, suffix = ''):
        if self.retrain:
            self.model.load_state_dict(torch.load('data/enhanced_models/' + self.model.model_name + '_' + self.name + suffix + '.pt'))
        else:
            self.model.load_raw_state()
